
import dbm
import time
import sys

def getDbmTable(name, mode='r', MAX_ATTEMPTS=1, SLEEP_INTERVAL=10):
    attempts = 0
    while attempts < MAX_ATTEMPTS:
        if attempts > 0:
            print >> sys.stderr, "# Waiting for resource", name
            time.sleep(SLEEP_INTERVAL)
        attempts += 1
        try:
            db = dbm.open(name, mode)
            return db
        except Exception, e:
            pass
    return None

def createExpressionDB(name, expressionFile):
    db = getDbmTable(name, 'w')
    if db==None:
        raise Exception("couldn't create database with name", name)
    for key in db.keys()[:]:
        del db[key]
    for line in open(expressionFile):
        cols = line.strip().split("\t")
        if len(cols)<=1 or line.startswith("#"): continue
        if cols[1]=="Exon" or cols[1]=="Intron":
            key = "%s\t%s\t%s" % (cols[0], cols[1], cols[2])
            val = "%s\t%s" % (cols[8], cols[11])
            db[key] = val
        elif cols[1]=="GeneExpression":
            key = "%s\t%s" % (cols[0], cols[1])
            val = "%s\t%s" % (cols[9], cols[12])
            db[key] = val
    db.close()

def readGeneExpression(name, gene, exon=None):
    db = getDbmTable(name, 'r')
    if db==None:
        raise Exception("couldn't open database with name", name)
    if exon!=None:
        key = "%s\t%s\t%s" % (gene, "Exon", exon)
    else:
        key = "%s\t%s" % (gene, "GeneExpression")
    if db.has_key(key):
        val = db[key]
    else:
        val = "NA"
    db.close()
    return val

def readMultiple(name, queries):
    db = getDbmTable(name, 'r')
    if db==None:
        raise Exception("couldn't open database with name", name)
    res = []
    for (gene, exon) in queries:
        if exon!=None:
            key = "%s\t%s\t%s" % (gene, "Exon", exon)
        else:
            key = "%s\t%s" % (gene, "GeneExpression")
        if db.has_key(key):
            res.append(db[key])
        else:
            res.append("NA")
    db.close()
    return res
    
if __name__=="__main__":
    if sys.argv[1]=="-c":
        createExpressionDB(sys.argv[2], sys.argv[3])
    elif len(sys.argv)>3 and '-' in sys.argv[3]:
        gene = sys.argv[2]
        exonRange = sys.argv[3].split("-")
        query = "Gene " + gene + " Exons " + sys.argv[3]
        minExon = int(exonRange[0])
        maxExon = int(exonRange[1])
        queries = [(gene, exon) for exon in range(minExon, maxExon+1)]
        res = readMultiple(sys.argv[1], queries)
        print >> sys.stderr, "# Multiple results for query '" + query + "' = " + str(res)
    else:
        exon = None
        query = "Gene " + sys.argv[2]
        if len(sys.argv)>3:
            exon = sys.argv[3]
            query = "Gene " + sys.argv[2] + " Exon " + sys.argv[3]
        val = readGeneExpression(sys.argv[1], sys.argv[2], exon)
        print >> sys.stderr, "# Result for query '" + query + "' = " + val
